from typing import Tuple

from torchvision.transforms import Compose

from .custom import AbstractDataset


# Decorator pattern


class TransformFirstDataset(AbstractDataset):

    def __init__(self, dataset: AbstractDataset, transform: Compose):
        super().__init__(dataset.root, dataset.partition, lazy_load=True)
        self.dataset = dataset
        self.transform = transform

    def load(self):
        self.dataset.load()

    def __len__(self) -> int:
        return len(self.dataset)

    def __getitem__(self, index: int) -> Tuple:
        data = self.dataset[index]
        transformed = []

        for i, item in enumerate(data):
            if i == 0:
                transformed.append(self.transform(item))
            else:
                transformed.append(item)

        return tuple(transformed)
